
##---------------------------------------------------------------
## Required libraries
##---------------------------------------------------------------

#----- Parallel computing
library(doParallel)

#----- Make clusters based on the number of CPU cores
cl <- makeCluster(detectCores())
registerDoParallel(cl)
getDoParWorkers()

#----- Support parallel excution
library(foreach)

#----- Stan (MCMC sampling) for R
library(rstan)

##---------------------------------------------------------------
## Extract MCMC samples from Stan outputs
##---------------------------------------------------------------

#----- Load MCMC samples and Data
load("MCMCsample.RData")
load("Master.RData")

#------ Set Treatment (TRT), Outcome (OUT), Mediators (M) and Covariates (X)
Data <- Master
OUT <- Data$PM.2.5
TRT <- Data$SO2.SC
M <- cbind(Data$SO2_Annual, Data$NOx_Annual, Data$CO2_Annual)
X <- cbind(Data$S_n_CR, Data$NumNOxControls, Data$Heat_Input/100000, Data$Barometric_Pressure, Data$Temperature,  Data$PctCapacity, Data$sulfur_Content, Data$Phase2_Indicator, Data$Operating_Time/1000)


dim.cov <- dim(X)[2] #<--------- Num. of Covariates

#------ Variables by treatments
x0 <- X[which(TRT==0),]
x1 <- X[which(TRT==1),]

y0 <- OUT[which(TRT==0)]
y1 <- OUT[which(TRT==1)]

m0 <- log(M[which(TRT==0),])
m1 <- log(M[which(TRT==1),])

n0 <- dim(x0)[1]
n1 <- dim(x1)[1]


#----- Extract MCMC samples from each marginal distribution

data.y0 <- extract(fit.y0)
data.y1 <- extract(fit.y1)
data.m10 <- extract(fit.m1_0)
data.m11 <- extract(fit.m1_1)
data.m20 <- extract(fit.m2_0)
data.m21 <- extract(fit.m2_1)
data.m30 <- extract(fit.m3_0)
data.m31 <- extract(fit.m3_1)

#----- Parameters of the models under Z=1

gamma10 <- data.y1$gamma0
gamma11 <- data.y1$gamma1
w1 <- data.y1$W
beta1_10 <- data.m11$gamma0
beta1_11 <- data.m11$gamma1
beta2_10 <- data.m21$gamma0
beta2_11 <- data.m21$gamma1
beta3_10 <- data.m31$gamma0
beta3_11 <- data.m31$gamma1
s1_1 <- data.m11$W
s2_1 <- data.m21$W
s3_1 <- data.m31$W
psi1 <- data.y1$psi
psi1_1 <- data.m11$psi
psi2_1 <- data.m21$psi
psi3_1 <- data.m31$psi


#----- Parameters of the models under Z=0

gamma00 <- data.y0$gamma0
gamma01 <- data.y0$gamma1
w0 <- data.y0$W
beta1_00 <- data.m10$gamma0
beta1_01 <- data.m10$gamma1
beta2_00 <- data.m20$gamma0
beta2_01 <- data.m20$gamma1
beta3_00 <- data.m30$gamma0
beta3_01 <- data.m30$gamma1
s1_0 <- data.m10$W
s2_0 <- data.m20$W
s3_0 <- data.m30$W
psi0 <- data.y0$psi
psi1_0 <- data.m10$psi
psi2_0 <- data.m20$psi
psi3_0 <- data.m30$psi


#----- Setting Thinning and Burn-in's
Thin <- 5
Burn0 <- 18000  # Extra burn-in periods for Models under Z=0
Burn1 <- 18000   # Extra burn-in periods for Models under Z=1

#----- Set the number of post-processing steps, N
n.iter <- 400


##---------------------------------------------------------------
## 'Main' function on each cluster (parallel)
##---------------------------------------------------------------

main<-function(temp){
    
    #----- Require to fit multivariate normal distributions
    library(mnormt)

    joint0 <- cbind(y0, m0);  joint1 <- cbind(y1, m1)
    #----- Estimating condtional correlations
    fit0 <- lm(joint0 ~ x0)
    fit1 <- lm(joint1 ~ x1)

    cov0 <- 1 / n0 * t(joint0 - predict(fit0, newdata=data.frame(x0))) %*% (joint0-predict(fit0, newdata = data.frame(x0)))
    cov1 <- 1 / n1 * t(joint1 - predict(fit1, newdata=data.frame(x1))) %*% (joint1-predict(fit1, newdata = data.frame(x1)))

    #----- Under Normality, the following relatioship between
    #----- Spearman and Pearson correlations holds (Kruskal, 1958)
    cor0 <- 6 / pi * asin(matrix(apply(expand.grid(1:4,1:4), 1, function(x) cov0[x[1],x[2]] / sqrt(cov0[x[1],x[1]] * cov0[x[2],x[2]])),4,4)/2)
    cor1 <- 6 / pi * asin(matrix(apply(expand.grid(1:4,1:4), 1, function(x) cov1[x[1],x[2]] / sqrt(cov1[x[1],x[1]]*cov1[x[2],x[2]])),4,4)/2)


    #----------------- Sensitivity parameters
    sens1 <- function(j,k) 2 * min(abs(cor0[j+1,k+1]), abs(cor1[j+1,k+1])) / abs(cor0[j+1,k+1] + cor1[j+1,k+1])
    rm <- function(j,k,rho1) (cor0[j+1,k+1] + cor1[j+1,k+1]) / 2 * rho1

    sens2 <- function(k) 2 * min(abs(cor0[1,k+1]), abs(cor1[1,k+1])) / abs(cor0[1,k+1]+ cor1[1,k+1])
    ry <- function(k,rho2) (cor0[1,k+1]+ cor1[1,k+1]) / 2 * rho2
    rho1 <- runif(1, 0, min(apply(expand.grid(1:3,1:3),1, function(x) sens1(x[1],x[2]))))
    rho2 <- runif(1, 0, min(sens2(1),sens2(2), sens2(3)))
    
    #----- Construct the correlation matrix
    cor.part <- matrix(apply(expand.grid(1:3,1:3), 1, function(x) rm(x[1], x[2], rho1)), 3, 3)
    cor.y1m0 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
    cor.y0m1 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
    
    COR1 <- 2 * sin(pi * (1/6) *rbind(cbind(cor1,rbind(cor.y1m0, cor.part)), cbind(t(rbind(cor.y1m0, cor.part)), cor0[2:4,2:4] ) ))
    COR0 <- 2 * sin(pi * (1/6) *rbind(cbind(cor0,rbind(cor.y0m1, cor.part)), cbind(t(rbind(cor.y0m1, cor.part)), cor1[2:4,2:4] ) ))
    
    #----- Check positive definiteness
    S0 <- attr(chol(COR0, pivot=TRUE),"rank");     S1 <- attr(chol(COR1, pivot=TRUE),"rank")
    while (S0 !=7 | S1 !=7){
        rho1 <- runif(1, 0, min(apply(expand.grid(1:3,1:3),1, function(x) sens1(x[1],x[2]))))
        rho2 <- runif(1, 0, min(sens2(1),sens2(2), sens2(3)))
        cor.part <- matrix(apply(expand.grid(1:3,1:3), 1, function(x) rm(x[1], x[2], rho1)), 3, 3)
        cor.y1m0 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
        cor.y0m1 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
        COR1 <- 2 * sin(pi * (1/6) *rbind(cbind(cor1,rbind(cor.y1m0, cor.part)), cbind(t(rbind(cor.y1m0, cor.part)), cor0[2:4,2:4] ) ))
        COR0 <- 2 * sin(pi * (1/6) *rbind(cbind(cor0,rbind(cor.y0m1, cor.part)), cbind(t(rbind(cor.y0m1, cor.part)), cor1[2:4,2:4] ) ))
        S0 <- attr(chol(COR0, pivot=TRUE),"rank");     S1 <- attr(chol(COR1, pivot=TRUE),"rank")
    }

    #----- Index of iterations (on parallel)
    j <- temp

    #----- Index of posterior samples
    index0 <- Thin * j + Burn0; index1 <- Thin * j + Burn1


    #----- The number of covariates samples (with replacement) for the empirical distribution
    size <- 30000
    s.covariate <- data.frame(X[sample(seq(1,dim(X)[1]), size = size, replace = TRUE),])  #----- Covariates samples


    #----- Weighting parameters of mixtures
    pi.y0 <- pmax(psi0[index0,], 0)
    pi.y1 <- pmax(psi1[index1,], 0)
    pi.m10 <- pmax(psi1_0[index0,], 0)
    pi.m20 <- pmax(psi2_0[index0,], 0)
    pi.m30 <- pmax(psi3_0[index0,], 0)
    pi.m11 <- pmax(psi1_1[index1,], 0)
    pi.m21 <- pmax(psi2_1[index1,], 0)
    pi.m31 <- pmax(psi3_1[index1,], 0)

    ##----- Sampling from the joint distribution of [Y(1;M(1,1,1)), M(0,0,0), M(1,1,1)]
    f1.M0M1 <- function(rho1, k){
        Cor01 <- COR1
    
        #----- Random samples from the Copula and take the standard normal CDF on them
        F <- pnorm(rmnorm(1, mean=rep(0, 7), Cor01), 0, 1)
        X.temp <- s.covariate[k,]
    
        #----- Draw a cluster for each marginal distribution
        clus.y1 <- which(rmultinom(1, 1, pi.y1) == 1)
        clus.m10 <- which(rmultinom(1, 1, pi.m10) == 1)
        clus.m20 <- which(rmultinom(1, 1, pi.m20) == 1)
        clus.m30 <- which(rmultinom(1, 1, pi.m30) == 1)
        clus.m11 <- which(rmultinom(1, 1, pi.m11) == 1)
        clus.m21 <- which(rmultinom(1, 1, pi.m21) == 1)
        clus.m31 <- which(rmultinom(1, 1, pi.m31) == 1)
    
        #----- Samples from marginal distributions
        s.y1 <- qnorm(F[1], mean = gamma10[index1, clus.y1] + gamma11[index1, ] %*% t(X.temp), sd = 1 / sqrt(w1[index1, clus.y1]))
        s.m1_0 <- qnorm(F[5], mean = beta1_00[index0, clus.m10] + beta1_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s1_0[index0, clus.m10]))
        s.m2_0 <- qnorm(F[6], mean = beta2_00[index0, clus.m20] + beta2_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s2_0[index0, clus.m20]))
        s.m3_0 <- qnorm(F[7], mean = beta3_00[index0, clus.m30] + beta3_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s3_0[index0, clus.m30]))
    
        s.m1_1 <- qnorm(F[2], mean = beta1_10[index1, clus.m11] + beta1_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s1_1[index1, clus.m11]))
        s.m2_1 <- qnorm(F[3], mean = beta2_10[index1, clus.m21] + beta2_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s2_1[index1, clus.m21]))
        s.m3_1 <- qnorm(F[4], mean = beta3_10[index1, clus.m31] + beta3_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s3_1[index1, clus.m31]))
        return(c(s.y1, s.m1_0, s.m2_0, s.m3_0, s.m1_1, s.m2_1, s.m3_1))
    }

    ##----- Sampling from the joint distribution of [Y(0;M(0,0,0)), M(0,0,0), M(1,1,1)]
    f0.M0M1 <- function(rho1, k){
        Cor01 <- COR0
    
        #----- Random samples from the Copula and take the standard normal CDF on them
        F <- pnorm(rmnorm(1, mean = rep(0, 7), Cor01), 0, 1)
        X.temp <- s.covariate[k, ]
    
        #----- Draw a cluster for each marginal distribution
        clus.y0 <- which(rmultinom(1, 1, pi.y0) == 1)
        clus.m10 <- which(rmultinom(1, 1, pi.m10) == 1)
        clus.m20 <- which(rmultinom(1, 1, pi.m20) == 1)
        clus.m30 <- which(rmultinom(1, 1, pi.m30) == 1)
        clus.m11 <- which(rmultinom(1, 1, pi.m11) == 1)
        clus.m21 <- which(rmultinom(1, 1, pi.m21) == 1)
        clus.m31 <- which(rmultinom(1, 1, pi.m31) == 1)
    
        #----- Samples from marginal distributions
        s.y0 <- qnorm(F[1], mean = gamma00[index0, clus.y0] + gamma01[index0, ] %*% t(X.temp), sd = 1 / sqrt(w0[index0, clus.y0]))
        s.m1_0 <- qnorm(F[2], mean = beta1_00[index0, clus.m10] + beta1_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s1_0[index0, clus.m10]))
        s.m2_0 <- qnorm(F[3], mean = beta2_00[index0, clus.m20] + beta2_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s2_0[index0, clus.m20]))
        s.m3_0 <- qnorm(F[4], mean = beta3_00[index0, clus.m30] + beta3_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s3_0[index0, clus.m30]))
    
        s.m1_1 <- qnorm(F[5], mean = beta1_10[index1, clus.m11] + beta1_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s1_1[index1, clus.m11]))
        s.m2_1 <- qnorm(F[6], mean = beta2_10[index1, clus.m21] + beta2_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s2_1[index1, clus.m21]))
        s.m3_1 <- qnorm(F[7], mean = beta3_10[index1, clus.m31] + beta3_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s3_1[index1, clus.m31]))
        return(c(s.y0, s.m1_0, s.m2_0, s.m3_0, s.m1_1, s.m2_1, s.m3_1))
    }


    #----- Estimating Y1 and Y0 conditional on all mediators
    y1.sample <- sapply(seq(1, dim(s.covariate)[1], by=1), function(x) f1.M0M1(rho1, x ))
    y0.sample <- sapply(seq(1, dim(s.covariate)[1], by=1), function(x) f0.M0M1(rho1, x ))


    #----- Thresholds (determined by sd's of the estimated individua-level causal effect on every emissions)
    c <- NULL
    c[1] <- 1.386*0.5
    c[2] <- 1.103*0.5
    c[3] <- 1.015*0.5


    #----- Define Strata
    for(k in 1:3){
        eval(parse(text=paste("abs1",k,"<- which(abs(y1.sample[",k,"+1,]-y1.sample[",k,"+4,]) <= c[",k,"])", sep="")))
        eval(parse(text=paste("abs0",k,"<- which(abs(y0.sample[",k,"+1,]-y0.sample[",k,"+4,]) <= c[",k,"])", sep="")))
        eval(parse(text=paste("pos1",k,"<- which(y1.sample[",k,"+1,]-y1.sample[",k,"+4,] >= c[",k,"])", sep="")))
        eval(parse(text=paste("pos0",k,"<- which(y0.sample[",k,"+1,]-y0.sample[",k,"+4,] >= c[",k,"])", sep="")))
        eval(parse(text=paste("neg1",k,"<- which(y1.sample[",k,"+1,]-y1.sample[",k,"+4,] <= -c[",k,"])", sep="")))
        eval(parse(text=paste("neg0",k,"<- which(y0.sample[",k,"+1,]-y0.sample[",k,"+4,] <= -c[",k,"])", sep="")))
    }
    
    #----- Estimating Principal Causal Effects
    EAE1 <- NULL
    EAE2 <- NULL
    EDE <- NULL
    Pr <- NULL

    for(k in 1:3){
        eval(parse(text=paste("mean.y1 <- mean(y1.sample[1,abs1",k,"]); mean.y0 <- mean(y0.sample[1,abs0",k,"]); EDE[",k,"] <- mean.y1-mean.y0", sep="")))
        eval(parse(text=paste("mean.y1 <- mean(y1.sample[1,pos1",k,"]); mean.y0 <- mean(y0.sample[1,pos0",k,"]); EAE1[",k,"] <- mean.y1-mean.y0", sep="")))
        eval(parse(text=paste("mean.y1 <- mean(y1.sample[1,neg1",k,"]); mean.y0 <- mean(y0.sample[1,neg0",k,"]); EAE2[",k,"] <- mean.y1-mean.y0", sep="")))
        
        Pr[3*k-2] <- eval(parse(text=paste("length(abs1",k,")/size", sep="")))
        Pr[3*k-1] <- eval(parse(text=paste("length(pos1",k,")/size", sep="")))
        Pr[3*k] <- eval(parse(text=paste("length(neg1",k,")/size", sep="")))
    }

    for(k in 1:3){
        index<-matrix(c(1,2,1,3,2,3),byrow=T, nrow=3)
        eval(parse(text=paste("mean.y1 <- mean(y1.sample[1,intersect(abs1",index[k,1],", abs1",index[k,2],")]); mean.y0 <- mean(y0.sample[1, intersect(abs0",index[k,1],", abs0",index[k,2],")]); EDE[",k+3,"] <- mean.y1-mean.y0", sep="")))
        eval(parse(text=paste("mean.y1 <- mean(y1.sample[1,intersect(pos1",index[k,1],", pos1",index[k,2],")]); mean.y0 <- mean(y0.sample[1, intersect(pos0",index[k,1],", pos0",index[k,2],")]); EAE1[",k+3,"] <- mean.y1-mean.y0", sep="")))
        eval(parse(text=paste("mean.y1 <- mean(y1.sample[1,intersect(neg1",index[k,1],", neg1",index[k,2],")]); mean.y0 <- mean(y0.sample[1, intersect(neg0",index[k,1],", neg0",index[k,2],")]); EAE2[",k+3,"] <- mean.y1-mean.y0", sep="")))
    
        eval(parse(text=paste("Pr[",3*k+7,"] <- length(intersect(abs1",index[k,1],",abs1",index[k,2],")) /size", sep="")))
        eval(parse(text=paste("Pr[",3*k+8,"] <- length(intersect(pos1",index[k,1],",pos1",index[k,2],")) /size", sep="")))
        eval(parse(text=paste("Pr[",3*k+9,"] <- length(intersect(neg1",index[k,1],",neg1",index[k,2],")) /size", sep="")))
    }

    mean.y1 <- mean(y1.sample[1, Reduce(function(x, y) intersect(x,y),list(abs11, abs12, abs13))])
    mean.y0 <- mean(y0.sample[1, Reduce(function(x, y) intersect(x,y),list(abs01, abs02, abs03))])
    EDE[7] <- mean.y1-mean.y0

    mean.y1 <- mean(y1.sample[1, Reduce(function(x, y) intersect(x,y),list(pos11, pos12, pos13))])
    mean.y0 <- mean(y0.sample[1, Reduce(function(x, y) intersect(x,y),list(pos01, pos02, pos03))])
    EAE1[7] <- mean.y1-mean.y0

    mean.y1 <- mean(y1.sample[1, Reduce(function(x, y) intersect(x,y),list(neg11, neg12, neg13))])
    mean.y0 <- mean(y1.sample[1, Reduce(function(x, y) intersect(x,y),list(neg01, neg02, neg03))])
    EAE2[7] <- mean.y1-mean.y0

    Pr[19] <- length(Reduce(function(x, y) intersect(x,y),list(abs11, abs12, abs13))) / size
    Pr[20] <- length(Reduce(function(x, y) intersect(x,y),list(pos11, pos12, pos13))) / size
    Pr[21] <- length(Reduce(function(x, y) intersect(x,y),list(neg11, neg12, neg13))) / size

return(c(EAE1[1:7], EDE[1:7], EAE2[1:7], Pr[1:21]))

}

result<-foreach(temp = 1:n.iter, .combine = rbind) %dopar% main(temp)

save(result, file="result_PS.RData")
stopCluster(cl)
